import os
import cv2
import torch
import numpy as np
from ultralytics import YOLO
import torch.nn as nn

class PersonDetector:
    def __init__(self, model_path='train/weights/best.pt', conf_thres=0.5):
        self.model = YOLO(model_path)
        self.conf_thres = conf_thres
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        
        # 修复 'Upsample' object has no attribute 'recompute_scale_factor' 错误
        for m in self.model.modules():
            if isinstance(m, nn.Upsample):
                m.recompute_scale_factor = None
        
    def detect(self, frame, classes=0):
        try:
            results = self.model(frame, conf=self.conf_thres, classes=classes)
            
            detections = []
            for r in results:
                try:
                    boxes = r.boxes.cpu().numpy()
                    for box in boxes:
                        try:
                            x1, y1, x2, y2 = box.xyxy[0]
                            conf = box.conf[0]
                            cls = box.cls[0]
                            detections.append([x1, y1, x2, y2, conf, cls])
                        except Exception as box_error:
                            print(f"处理检测框时出错: {box_error}")
                            continue
                except Exception as boxes_error:
                    print(f"处理检测结果时出错: {boxes_error}")
                    continue
            
            # 确保返回 numpy 数组
            if not detections:
                print("未检测到任何目标")
                return np.empty((0, 6))
                
            return np.array(detections)
        except Exception as e:
            print(f"检测过程中出错: {e}")
            # 发生错误时返回空数组
            return np.empty((0, 6)) 